﻿using System.Collections.Generic;
using System.Linq;

namespace CodeContrib.Data.Entity
{
	public static class TreeExtensions
	{
		public static TTreeNode ParentOf<TTreeNode>(this IEnumerable<TTreeNode> treeNodes, TreeNode treeNode) where TTreeNode : TreeNode, new()
		{
			return treeNodes.SingleOrDefault(tn => tn.Left < treeNode.Left && tn.Right > treeNode.Right && tn.Depth == treeNode.Depth - 1);
		}

		public static TTreeNode Root<TTreeNode>(this IEnumerable<TTreeNode> treeNodes) where TTreeNode : TreeNode, new()
		{
			return treeNodes.SingleOrDefault(tn => tn.Left == 0);
		}

		public static IEnumerable<TreeHierarchy<TTreeNode>> ToHierarchy<TTreeNode>(this IEnumerable<TTreeNode> treeNodes, TTreeNode parentTreeNode) where TTreeNode : TreeNode, new()
		{
			treeNodes = treeNodes as TTreeNode[] ?? treeNodes.ToArray();

			foreach (var treeNode in treeNodes.WhereChildrenOf(parentTreeNode))
			{
				yield return new TreeHierarchy<TTreeNode>
					             {
									 Children = ToHierarchy(treeNodes.WhereDescendantsOf(treeNode), treeNode),
						             TreeNode = treeNode
					             };
			}
		}

		public static IEnumerable<TTreeNode> WhereAncestorsOf<TTreeNode>(this IEnumerable<TTreeNode> treeNodes, TreeNode treeNode) where TTreeNode : TreeNode, new()
		{
			return treeNodes.Where(tn => tn.Left > treeNode.Left && tn.Right < treeNode.Right && tn.Depth < treeNode.Depth);
		}

		public static IEnumerable<TTreeNode> WhereChildrenOf<TTreeNode>(this IEnumerable<TTreeNode> treeNodes, TreeNode treeNode) where TTreeNode : TreeNode, new()
		{
			return treeNodes.Where(tn => tn.Left > treeNode.Left && tn.Right < treeNode.Right && tn.Depth == treeNode.Depth + 1);
		}

		public static IEnumerable<TTreeNode> WhereDescendantsOf<TTreeNode>(this IEnumerable<TTreeNode> treeNodes, TreeNode treeNode) where TTreeNode : TreeNode, new()
		{
			return treeNodes.Where(tn => tn.Left > treeNode.Left && tn.Right < treeNode.Right && tn.Depth > treeNode.Depth);
		}

		public static IEnumerable<TTreeNode> WhereLastDescendants<TTreeNode>(this IEnumerable<TTreeNode> treeNodes) where TTreeNode : TreeNode, new()
		{
			return treeNodes.Where(tn => tn.Right - tn.Left == 1);
		}
	}
}